#include <iostream>
#include <random>
#include <cstring>
#include <chrono>
#include <glm/mat4x4.hpp>
#include <glm/gtx/transform.hpp>
#include "VK_UniformBuffer.h"
#include "VK_Context.h"
#include "VK_Pipeline.h"
#include "VK_Image.h"
#include "VK_Texture.h"
#include "VK_DynamicState.h"
#include "VK_Window.h"
#include "VK_CommandPool.h"
#include "../Node.h"
#include "../ObjectFactory.h"

using namespace std;

const std::vector<float> vertices = {
    -1.0f, -1.0f,  1.0f, 1.0f, 0.0f, 0.0f,
    1.0f, -1.0f,  1.0f, 0.0f, 1.0f, 0.0f,
    1.0f,  1.0f,  1.0f, 0.0f, 0.0f, 1.0f,
    -1.0f,  1.0f,  1.0f, 0.0f, 0.0f, 0.0f,
    -1.0f, -1.0f, -1.0f, 1.0f, 0.0f, 0.0f,
    1.0f, -1.0f, -1.0f, 0.0f, 1.0f, 0.0f,
    1.0f,  1.0f, -1.0f, 0.0f, 0.0f, 1.0f,
    -1.0f,  1.0f, -1.0f, 0.0f, 0.0f, 0.0f,
};

const std::vector<uint32_t> indices = {
    0,1,2, 2,3,0, 1,5,6, 6,2,1, 7,6,5, 5,4,7, 4,0,3, 3,7,4, 4,5,1, 1,0,4, 3,2,6, 6,7,3,
};

#define OBJECT_INSTANCES 125

struct UBO
{
    glm::mat4 projection;
    glm::mat4 view;
};


class dynamicuniform : public Node
{
public:
    dynamicuniform()
    {
        nodeName = "dynamicuniform";
    }
protected:
    RunningState run()override
    {
        auto window = VK_Window::getInstance()->getWindow();
        auto title = getNodeTitle();
        glfwSetWindowTitle(window, title.data());

        context = createVkContext();
        VK_Context::VK_Config vkConfig;
        vkConfig.usingDepthAttachment = true;
        bool ok = context->initVulkanDevice(vkConfig);
        if (!ok)
            return RunningState_Failed;

        auto property = context->getPhysicalDeviceProperties();
        size_t minUboAlignment = property.limits.minUniformBufferOffsetAlignment;
        dynamicAlignment = sizeof(glm::mat4);
        if (minUboAlignment > 0) {
            dynamicAlignment = (dynamicAlignment + minUboAlignment - 1) & ~(minUboAlignment - 1);
        }
        dynamicUniformSize = OBJECT_INSTANCES * dynamicAlignment;

        auto shaderSet = context->createShaderSet();
        shaderSet->addShader("../shader/dynamic-uniform/vert.spv", VK_SHADER_STAGE_VERTEX_BIT);
        shaderSet->addShader("../shader/dynamic-uniform/frag.spv", VK_SHADER_STAGE_FRAGMENT_BIT);

        shaderSet->appendVertexAttributeDescription(0, sizeof (float) * 3, VK_FORMAT_R32G32B32_SFLOAT, 0);
        shaderSet->appendVertexAttributeDescription(1, sizeof (float) * 3, VK_FORMAT_R32G32B32_SFLOAT,
                                                    sizeof(float) * 3);

        shaderSet->appendVertexInputBindingDescription(6 * sizeof(float), 0, VK_VERTEX_INPUT_RATE_VERTEX);

        VkDescriptorSetLayoutBinding uniformBinding = VK_ShaderSet::createDescriptorSetLayoutBinding(0,
            VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, VK_SHADER_STAGE_VERTEX_BIT);
        shaderSet->addDescriptorSetLayoutBinding(uniformBinding);

        VkDescriptorSetLayoutBinding dynamincuUniformBinding = VK_ShaderSet::createDescriptorSetLayoutBinding(1,
            VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, VK_SHADER_STAGE_VERTEX_BIT);
        shaderSet->addDescriptorSetLayoutBinding(dynamincuUniformBinding);

        if (!shaderSet->isValid()) {
            std::cerr << "invalid shaderSet" << std::endl;
            shaderSet->release();
            context->release();
            return RunningState_Failed;
        }

        auto ubo = shaderSet->addUniformBuffer(0, sizeof(UBO));
        ubo->setWriteDataCallback([&](char *&data, uint32_t size)->uint32_t {return updateUniformBufferData(data, size);});

        auto dynamicUbo = shaderSet->addDynamicUniformBuffer(1, dynamicAlignment, OBJECT_INSTANCES);
        dynamicUbo->setWriteDataCallback([&](char *&data, uint32_t size)->uint32_t{return updateDynamicUniformBufferData(data, size);});

        std::default_random_engine rndEngine((unsigned)time(nullptr));
        std::normal_distribution<float> rndDist(-1.0f, 1.0f);
        for (uint32_t i = 0; i < OBJECT_INSTANCES; i++) {
            rotations[i] = glm::vec3(rndDist(rndEngine), rndDist(rndEngine), rndDist(rndEngine)) * 2.0f * (float)M_PI;
            rotationSpeeds[i] = glm::vec3(rndDist(rndEngine), rndDist(rndEngine), rndDist(rndEngine));
        }

        int32_t width, height;
        glfwGetFramebufferSize(window, &width, &height);

        context->initVulkanContext();

        auto pipeline = context->createPipeline(shaderSet);
        pipeline->create();
        auto buffer = context->createVertexBuffer(vertices, 6, indices);
        pipeline->addRenderBuffer(buffer);

        context->createCommandBuffers();

        auto result = context->run();
        setAverageFps(context->getFps());
        context->release();
        return result;
    }

    uint32_t updateDynamicUniformBufferData(char *&data, uint32_t size)
    {
        animationTime += context->getFrameTime();
        if (animationTime <= 1.0f / 60.0f) {
            return dynamicUniformSize;
        }
        uint32_t dim = static_cast<uint32_t>(pow(OBJECT_INSTANCES, (1.0f / 3.0f)));
        glm::vec3 offset(5.0f);

        for (uint32_t x = 0; x < dim; x++)
        {
            for (uint32_t y = 0; y < dim; y++)
            {
                for (uint32_t z = 0; z < dim; z++)
                {
                    uint32_t index = x * dim * dim + y * dim + z;
                    rotations[index] += animationTime * rotationSpeeds[index];

                    glm::vec3 pos = glm::vec3(-((dim * offset.x) / 2.0f) + offset.x / 2.0f + x * offset.x, -((dim * offset.y) / 2.0f) + offset.y / 2.0f + y * offset.y, -((dim * offset.z) / 2.0f) + offset.z / 2.0f + z * offset.z);
                    glm::mat4 modelMat;
                    modelMat = glm::translate(glm::mat4(1.0f), pos);
                    modelMat = glm::rotate(modelMat, rotations[index].x, glm::vec3(1.0f, 1.0f, 0.0f));
                    modelMat = glm::rotate(modelMat, rotations[index].y, glm::vec3(0.0f, 1.0f, 0.0f));
                    modelMat = glm::rotate(modelMat, rotations[index].z, glm::vec3(0.0f, 0.0f, 1.0f));
                    memcpy(data + index * dynamicAlignment, &modelMat, dynamicAlignment);
                }
            }
        }
        animationTime = 0.0;
        return dynamicUniformSize;
    }

    uint32_t updateUniformBufferData(char *&data, uint32_t size)
    {
        assert(sizeof(UBO) == size);
        static auto startTime = std::chrono::high_resolution_clock::now();
        auto currentTime = std::chrono::high_resolution_clock::now();
        float time = std::chrono::duration<float, std::chrono::seconds::period>
                     (currentTime - startTime).count();
        glm::mat4 model = glm::identity<glm::mat4>();
        model *= glm::rotate(glm::mat4(1.0f), time * 0.15f * glm::radians(-90.0f), glm::vec3(0.0f, 0.0f, 1.0f));
        auto view = glm::lookAt(glm::vec3(50.0f, 0.0f, 0.0f), glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f, 0.0f, 1.0f));
        auto proj = glm::perspective(glm::radians(45.0f),
            context->getSwapChainExtent().width / (float)context->getSwapChainExtent().height, 0.1f, 100.0f);

        UBO ubo;
        ubo.projection = proj;
        ubo.view = view * model;
        memcpy(data, &ubo, sizeof(UBO));
        return sizeof(UBO);
    }

protected:
    VK_Context *context = nullptr;
    uint32_t dynamicAlignment = 0;
    uint32_t dynamicUniformSize = 0;
    glm::vec3 rotations[OBJECT_INSTANCES];
    glm::vec3 rotationSpeeds[OBJECT_INSTANCES];
    float animationTime = 0.0;
};

REGISTER_OBJECT(Node, dynamicuniform)
